Skip to content

Cpu memory graph break #3886

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Cpu memory graph break #3886

wants to merge 5 commits into from
+1,010 −0

Conversation

@cehongwang
Copy link
Collaborator

@cehongwang cehongwang commented Nov 4, 2025

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • This change requires a documentation update

Checklist:

  • My code follows the style guidelines of this project (You can use the linters)
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas and hacks
  • I have made corresponding changes to the documentation
  • I have added tests to verify my fix or my feature
  • New and existing unit tests pass locally with my changes
  • I have added the relevant labels to my PR in so that relevant reviewers are notified

@meta-cla meta-cla bot added the cla signed label Nov 4, 2025
@github-actions github-actions bot added component: tests Issues re: Tests component: lowering Issues re: The lowering / preprocessing passes component: conversion Issues re: Conversion stage component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Nov 4, 2025
@github-actions github-actions bot requested a review from peri044 November 4, 2025 20:05
github-actions[bot]
github-actions bot requested changes Nov 4, 2025
View reviewed changes
Copy link

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/_compiler.py	2025-11-04 20:05:23.825034+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/_compiler.py	2025-11-04 20:05:55.253944+00:00
@@ -876,15 +876,14 @@
    # This is done to release CPU memory.
    for attr in dir(gm):
        if attr.startswith("_frozen_param"):
            delattr(gm, attr)

-
-
    from torch_tensorrt.dynamo.conversion._ConverterRegistry import DYNAMO_CONVERTERS
+
    DYNAMO_CONVERTERS.disallowed_targets = set()
-    
+
    for name, _ in partitioned_module.named_children():
        submodule = getattr(partitioned_module, name)
        # filter on the GraphModule
        if not isinstance(submodule, torch.fx.graph_module.GraphModule):
            continue

narendasan
narendasan reviewed Nov 4, 2025
View reviewed changes
Copy link
Collaborator

@narendasan narendasan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you have a test case or something to demonstrate this feature?

@cehongwang cehongwang force-pushed the cpu-memory-graph-break branch from 7f0e504 to 18ccadf Compare November 5, 2025 22:03
@narendasan
Copy link
Collaborator

narendasan commented Nov 6, 2025

  1. We should think about using this tech for refit vs non refit
  2. Make refit apis work across graph breaks

@narendasan
Copy link
Collaborator

narendasan commented Nov 6, 2025

Improve usability by automating nn.Module -> atomic fx graph

@cehongwang cehongwang force-pushed the cpu-memory-graph-break branch from 18ccadf to f03ab2c Compare November 6, 2025 20:06
narendasan
narendasan reviewed Nov 7, 2025
View reviewed changes
narendasan
narendasan reviewed Nov 7, 2025
View reviewed changes
narendasan
narendasan reviewed Nov 7, 2025
View reviewed changes
narendasan
narendasan reviewed Nov 7, 2025
View reviewed changes
narendasan
narendasan reviewed Nov 7, 2025
View reviewed changes
py/torch_tensorrt/dynamo/_defaults.py Outdated
L2_LIMIT_FOR_TILING = -1
USE_DISTRIBUTED_MODE_TRACE = False
OFFLOAD_MODULE_TO_CPU = False
CPU_MEMORY_BUDGET = -1
Copy link
Collaborator

@narendasan narendasan Nov 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use an optional instead since this is not a TRT api we dont need -1 to mean let us decide

narendasan
narendasan reviewed Nov 7, 2025
View reviewed changes
narendasan
narendasan reviewed Nov 7, 2025
View reviewed changes
tests/py/dynamo/models/test_models.py
torch._dynamo.reset()


def compile_one(idx: int, ir: str):
Copy link
Collaborator

@narendasan narendasan Nov 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this test here?

narendasan
narendasan reviewed Nov 7, 2025
View reviewed changes
py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py Outdated

def size_of_subgraphs(self, subgraphs: List[Subgraph]) -> List[int]:
"""
This function calculates the size of the subgraph.
Copy link
Collaborator

@narendasan narendasan Nov 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you describe the algorithms here so we have reference for later?

@cehongwang cehongwang force-pushed the cpu-memory-graph-break branch 3 times, most recently from c2c7fce to 6f580b6 Compare November 13, 2025 19:11
@github-actions github-actions bot removed component: tests Issues re: Tests component: lowering Issues re: The lowering / preprocessing passes component: conversion Issues re: Conversion stage labels Nov 13, 2025
@cehongwang cehongwang force-pushed the cpu-memory-graph-break branch from 6f580b6 to b9fe0c1 Compare November 13, 2025 19:20
@github-actions github-actions bot added the component: tests Issues re: Tests label Nov 14, 2025
@cehongwang cehongwang force-pushed the cpu-memory-graph-break branch from e7cad5b to 99581e2 Compare November 14, 2025 19:13
narendasan
narendasan reviewed Nov 14, 2025
View reviewed changes
Copy link
Collaborator

@narendasan narendasan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Its looking good, just add a quick example in the examples folder and list it under contributor documentation for now

@cehongwang cehongwang force-pushed the cpu-memory-graph-break branch from 3b02ef9 to 290cc39 Compare November 14, 2025 20:19
zewenli98
zewenli98 requested changes Nov 14, 2025
View reviewed changes
@cehongwang cehongwang force-pushed the cpu-memory-graph-break branch from aba69ef to 756f827 Compare November 18, 2025 00:06
@cehongwang cehongwang linked an issue Nov 18, 2025 that may be closed by this pull request
✨[Feature] Resource aware Graph partitioner #3906
Open
@cehongwang cehongwang force-pushed the cpu-memory-graph-break branch from cb8480d to 2d6053e Compare November 19, 2025 21:20
narendasan
narendasan reviewed Nov 19, 2025
View reviewed changes
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@narendasan narendasan narendasan left review comments

@zewenli98 zewenli98 zewenli98 requested changes

@github-actions github-actions[bot] github-actions[bot] requested changes

@peri044 peri044 Awaiting requested review from peri044

Assignees

No one assigned

Labels

cla signed component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: tests Issues re: Tests

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

✨[Feature] Resource aware Graph partitioner

4 participants

@cehongwang @narendasan @zewenli98